Thanks for creating out the updated pull request @helpmefindaname and for pointing out the low accuracy @alanakbik!

I've looked at the low accuracies and we indeed had a bug with linking to the correct knowledge base identifiers. This part should be fixed now and A2M and alpha-2-macroglobulin point to the correct NCBI Gene identifier number 2.

@WangXII thanks for the update! Is there a way I can compute evaluation numbers in Flair? i.e. load a gold dataset, load the model, make predictions and evaluate? Could you post a snippet for this?

@sg-wbi can answer this the best. I think we have yet to update the evaluation script to the revised Flair API

@sg-wbi can you share an evaluation script? I would like to use it to test the models for accuracy before merging.

@sg-wbi can you share an evaluation script? I would like to use it to test the models for accuracy before merging.

Unfortunately this is not straightforward. When we developed this our efforts in the integration into flair stopped at the model level. This is becasue we had a very specific use case. At the current state our evaluation scripts require multiple preprocessing steps.

We (@mariosaenger, @WangXII) will get back to you asap with a workable solution (to your suggestions as well).

@alanakbik

Tests for accuracy

Here's the script that will give you accuracy results for 3 commonly used datasets.

from collections import defaultdict

from datasets import load_dataset

from flair.models import EntityMentionLinker
from flair.models.entity_mention_linking import BioSynEntityPreprocessor

ENTITY_TYPE_TO_MODEL = {
    "diseases": "dmis-lab/biosyn-sapbert-ncbi-disease",
    "chemical": "dmis-lab/biosyn-sapbert-bc5cdr-chemical",
    "genes": "dmis-lab/biosyn-sapbert-bc2gn",
}

ENTITY_TYPE_TO_DATASET = {"diseases": "ncbi_disease", "chemical": "bc5cdr", "genes": "gnormplus"}

ENITY_TYPE_TO_DICTIONARY = {"diseases": "ctd-diseases", "chemical": "ctd-chemicals", "genes": "ncbi-gene"}


def main():
    for entity_type, model in ENTITY_TYPE_TO_MODEL.items():
        ds_name = ENTITY_TYPE_TO_DATASET[entity_type]
        dictionary = ENITY_TYPE_TO_DICTIONARY[entity_type]

        ds = load_dataset(f"bigbio/{ds_name}", f"{ds_name}_bigbio_kb", trust_remote_code=True)
        print(f"Loaded corpus: `{ds_name}`")

        annotations = [a for d in ds["test"] for a in d["entities"]]
        if ds_name == "bc5cdr":
            annotations = [a for a in annotations if a["type"].lower() == entity_type]

        mention_to_uids = defaultdict(list)
        uid_to_link = {}
        for a in annotations:
            # skip mentions without normalization
            if len(a["normalized"]) == 0:
                continue

            if ds_name == "gnormplus":
                # no prefix for NCBI Gene
                uid_to_link[a["id"]] = [n["db_id"] for n in a["normalized"]]
            else:
                uid_to_link[a["id"]] = [":".join((n["db_name"], n["db_id"])) for n in a["normalized"]]

            for t in a["text"]:
                mention_to_uids[t].append(a["id"])

        linker = EntityMentionLinker.build(
            model,
            entity_type,
            dictionary_name_or_path=dictionary,
            hybrid_search=True,
            preprocessor=BioSynEntityPreprocessor(),
            batch_size=1024,
        )

        mentions = sorted(linker.preprocessor.process_mention(m) for m in set(mention_to_uids))
        results = linker.candidate_generator.search(entity_mentions=mentions, top_k=1)

        hits = 0
        total = 0
        for m, r in zip(mentions, results):
            for uid in mention_to_uids[m]:
                y_true = uid_to_link[uid]
                y_pred, _ = r[0]
                total += 1
                if y_pred in y_true:
                    hits += 1

        accuracy = round(hits / total * 100, 2)

        print(f"EVALUATION |  MODEL: `{model}`, CORPUS: `{ds_name}`,  ACCURACY@1: {accuracy}")


if __name__ == "__main__":
    main()

You should get the following results

EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-ncbi-disease`, CORPUS: `ncbi_disease`,  ACCURACY@1: 84.3
EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-bc5cdr-chemical`, CORPUS: `bc5cdr`,  ACCURACY@1: 93.85
EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-bc2gn`, CORPUS: `gnormplus`,  ACCURACY@1: 74.26

NOTE: the script reports accuracy on the gold mentions of the dataset (i.e. no NER) since this is how these models are commonly evaluated. Testing that EntityMentionLinker can correctly handle the output of a Classifier are in "test_biomedical_entity_linking.py"

NOTE: links to fetch models from the huggingface hub need to be updated. Let us know where we should put pre-trained models.

Suggestions

Is there some way of including the dictionary into the linker and preparing convenience functions to interpret the links?

Here we still have to find a nice solution. As you and @mariosaenger discussed last week - we may opt to build a distinct NENLabel class which inherits from Label and extends it by introducing two additional attributes: dictionary and canonical_name.

The label_type of all linking models could be "link" - not sure what is gained from having different label_type names for different biomedical NER classes.

For this, we adapted the implementation according to the approach done in the RelationClassifier model. Users are now able to specifically define on which label and entity types the linker model should be applied, e.g. by passing "ner" for all labels or {"ner": ["genes", "diseases"]} for restricting the linker to certain entity types for a given label type.

You could prepare MultitaskModels for each type (gene, disease, etc.) that combine the necessary NER tagger and linker, and correctly set the entity_tag_type.

Since this PR is already quite substantial, if that's ok with you, we would provide bundled models which correctly link the new NER and NEN models together when all code changes are done.

@sg-wbi

I am working on integrating you evaluation script into this PR.
While debugging, I think I found a bug:

        ....

-       mentions = sorted(linker.preprocessor.process_mention(m) for m in set(mention_to_uids))
-       results = linker.candidate_generator.search(entity_mentions=mentions, top_k=1)
+       mentions = sorted(mention_to_uids.keys())
+       preproc_mentions = [linker.preprocessor.process_mention(m) for m in mentions]
+       results = linker.candidate_generator.search(entity_mentions=preproc_mentions, top_k=1)

        hits = 0
        total = 0
        for m, r in zip(mentions, results):
        ....

Imo the before ignored all labels for mentiones that were preprocessed (e.g. the hard ones).

With the initial script I get the following results:

EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-ncbi-disease`, CORPUS: `ncbi_disease`,  ACCURACY@1: 83.91
EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-bc5cdr-chemical`, CORPUS: `bc5cdr`,  ACCURACY@1: 94.0
EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-bc2gn`, CORPUS: `gnormplus`,  ACCURACY@1: 69.31

after the change I get some worse results, with the exception of the genes corpus, that improves a lot:

EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-ncbi-disease`, CORPUS: `ncbi_disease`,  ACCURACY@1: 76.35
EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-bc5cdr-chemical`, CORPUS: `bc5cdr`,  ACCURACY@1: 84.62
EVALUATION |  MODEL: `dmis-lab/biosyn-sapbert-bc2gn`, CORPUS: `gnormplus`,  ACCURACY@1: 79.58

I also noticed, that some labels in the dataset do not exist in the dictionary and therefore cannot be predicted right.
I found the following upper bounds on the detection of the respective dictionary-dataset pairs:

EVALUATION |  NOT_IN_DATA: 81, MAX_ACCURACY@1: 91.56
EVALUATION |  NOT_IN_DATA: 271, MAX_ACCURACY@1: 95.06
EVALUATION |  NOT_IN_DATA: 227, MAX_ACCURACY@1: 92.93

Can you comfirm this?

I am working on integrating you evaluation script into this PR.

Thanks for taking care of this.

While debugging, I think I found a bug:
Imo the before ignored all labels for mentiones that were preprocessed (e.g. the hard ones).

Yes you are right. In my version only mentions which do not change after preprocessing were evaluated.

I also noticed, that some labels in the dataset do not exist in the dictionary and therefore cannot be predicted right.
Can you comfirm this?

Yes this is to be expected. Dictionary labels change over time (become obsolete/are merged). Some of these corpora were created 10 years ago

@alanakbik I have added the requested docstrings/comments.

@sg-wbi @helpmefindaname @mariosaenger thanks a lot for adding this major new feature to Flair!